{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Tce3stUlHN0L"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "tuOe1ymfHZPu"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qFdPvlXBOdUN"
      },
      "source": [
        "# 混合精度"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfBg1C5NB3X0"
      },
      "source": [
        "<table class=\"tfo-notebook-buttons\" align=\"left\">\n",
        "  <td><a target=\"_blank\" href=\"https://tensorflow.google.cn/guide/mixed_precision\" class=\"\"><img src=\"https://tensorflow.google.cn/images/tf_logo_32px.png\" class=\"\">在 TensorFlow.org 上查看</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/docs-l10n/blob/master/site/zh-cn/guide/mixed_precision.ipynb\" class=\"\"><img src=\"https://tensorflow.google.cn/images/colab_logo_32px.png\" class=\"\">在 Google Colab 中运行</a></td>\n",
        "  <td><a target=\"_blank\" href=\"https://github.com/tensorflow/docs-l10n/blob/master/site/zh-cn/guide/mixed_precision.ipynb\" class=\"\"><img src=\"https://tensorflow.google.cn/images/GitHub-Mark-32px.png\" class=\"\">在 GitHub 上查看源代码</a></td>\n",
        "  <td><a href=\"https://storage.googleapis.com/tensorflow_docs/docs-l10n/site/zh-cn/guide/mixed_precision.ipynb\" class=\"\"><img src=\"https://tensorflow.google.cn/images/download_logo_32px.png\" class=\"\">下载笔记本</a></td>\n",
        "</table>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xHxb-dlhMIzW"
      },
      "source": [
        "## 概述\n",
        "\n",
        "混合精度是指训练时在模型中同时使用 16 位和 32 位浮点类型，从而加快运行速度，减少内存使用的一种训练方法。通过让模型的某些部分保持使用 32 位类型以保持数值稳定性，可以缩短模型的单步用时，而在评估指标（如准确率）方面仍可以获得同等的训练效果。本文介绍如何使用实验性 Keras 混合精度 API 来加快模型速度。利用此 API 可以在现代 GPU 上将性能提高三倍以上，而在 TPU 上可以提高 60％。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AEDHaeuR7kLX"
      },
      "source": [
        "注：Keras 混合精度 API 目前是实验版本，可能会更改。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3vsYi_bv7gS_"
      },
      "source": [
        "如今，大多数模型使用 float32 dtype，这种数据类型占用 32 位内存。但是，还有两种精度较低的 dtype，即 float16 和 bfloat16，它们都是占用 16 位内存。现代加速器使用 16 位 dtype 执行运算的速度更快，因为它们有执行 16 位计算的专用硬件，并且从内存中读取 16 位 dtype 的速度也更快。\n",
        "\n",
        "NVIDIA GPU 使用 float16 执行运算的速度比使用 float32 快，而 TPU 使用 bfloat16 执行运算的速度也比使用 float32 快。因此，在这些设备上应尽可能使用精度较低的 dtype。但是，出于对数值的要求，为了让模型训练获得相同的质量，一些变量和计算仍需使用 float32。利用 Keras 混合精度 API，float16 或 bfloat16 可以与 float32 混合使用，从而既可以获得 float16/bfloat16 的性能优势，也可以获得 float32 的数值稳定性。\n",
        "\n",
        "注：在本指南中，术语“数值稳定性”是指使用较低精度的 dtype（而不是较高精度的 dtype）对模型质量的影响。如果使用 float16 或 bfloat16 执行运算，则与使用 float32 执行运算相比，使用这些较低精度的 dtype 会导致模型获得的计算准确率或其他指标相对较低，那么我们就说这种运算“数值不稳定”。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MUXex9ctTuDB"
      },
      "source": [
        "## 设置"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1Eh-iCRVBm0p"
      },
      "source": [
        "TensorFlow 2.1 中提供了 Keras 混合精度 API。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IqR2PQG4ZaZ0"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "\n",
        "from tensorflow import keras\n",
        "from tensorflow.keras import layers\n",
        "from tensorflow.keras.mixed_precision import experimental as mixed_precision"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "814VXqdh8Q0r"
      },
      "source": [
        "## 支持的硬件\n",
        "\n",
        "虽然混合精度在大多数硬件上都可以运行，但是在最新的 NVIDIA GPU 和 Cloud TPU 上才能加速模型。NVIDIA GPU 支持混合使用 float16 和 float32，而 TPU 则支持混合使用 bfloat16 和 float32。\n",
        "\n",
        "在 NVIDIA GPU 中，计算能力为 7.0 或更高的 GPU 可以获得混合精度的最大性能优势，因为这些型号具有称为 Tensor 核心的特殊硬件单元，可以加速 float16 矩阵乘法和卷积运算。旧款 GPU 使用混合精度无法实现数学运算性能优势，不过可以节省内存和带宽，因此也可以在一定程度上提高速度。您可以在 NVIDIA 的 [CUDA GPU 网页](https://developer.nvidia.com/cuda-gpus)上查询 GPU 的计算能力。可以最大程度从混合精度受益的 GPU 示例包括 RTX GPU、Titan V 和 V100。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "-q2hisD60F0_"
      },
      "source": [
        "注：如果在 Google Colab 中运行本指南中示例，则 GPU 运行时通常会连接 P100。P100 的计算能力为 6.0，预计速度提升不明显。\n",
        "\n",
        "您可以使用以下命令检查 GPU 类型。如果要使用此命令，必须安装 NVIDIA 驱动程序，否则会引发错误。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "j-Yzg_lfkoa_"
      },
      "outputs": [],
      "source": [
        "!nvidia-smi -L"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hu_pvZDN0El3"
      },
      "source": [
        "所有 Cloud TPU 均支持 bfloat16。\n",
        "\n",
        "即使在预计无法提升速度的 CPU 和旧款 GPU 上，混合精度 API 仍可以用于单元测试、调试，或试用 API​​。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HNOmvumB-orT"
      },
      "source": [
        "## 设置 dtype 策略"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "54ecYY2Hn16E"
      },
      "source": [
        "要在 Keras 中使用混合精度，您需要创建一条 `tf.keras.mixed_precision.experimental.Policy`，通常将其称为 *dtype 策略*。Dtype 策略可以指定将在其中运行的 dtype 层。在本指南中，您将从字符串 `'mixed_float16'` 构造策略，并将其设置为全局策略。这会导致随后创建的层使用 float16 和 float32 的混合精度。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "x3kElPVH-siO"
      },
      "outputs": [],
      "source": [
        "policy = mixed_precision.Policy('mixed_float16')\n",
        "mixed_precision.set_policy(policy)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oGAMaa0Ho3yk"
      },
      "source": [
        "该策略指定了层的两个重要方面：完成层的计算所使用的 dtype 和层变量的 dtype。上面的代码创建了一条 `mixed_float16` 策略（即通过将字符串 `'mixed_float16'` 传递给其构造函数而构建的 `mixed_precision.Policy` ）。凭借此策略，层可以使用 float16 计算和 float32 变量。计算使用 float16 来提高性能，而变量使用 float32 来确保数值稳定性。您可以直接在策略中查询这些属性。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GQRbYm4f8p-k"
      },
      "outputs": [],
      "source": [
        "print('Compute dtype: %s' % policy.compute_dtype)\n",
        "print('Variable dtype: %s' % policy.variable_dtype)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MOFEcna28o4T"
      },
      "source": [
        "如前所述，在计算能力至少为 7.0 的 NVIDIA GPU 上，`mixed_float16` 策略可以大幅提升性能。在其他 GPU 和 CPU 上，该策略也可以运行，但可能无法提升性能。对于 TPU，则应使用 `mixed_bfloat16` 策略。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cAHpt128tVpK"
      },
      "source": [
        "## 构建模型"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "nB6ujaR8qMAy"
      },
      "source": [
        "接下来，我们开始构建一个简单的模型。过小的模型往往无法获得混合精度的优势，因为 TensorFlow 运行时的开销通常占据大部分执行时间，导致 GPU 的性能提升几乎可以忽略不计。因此，如果使用 GPU，我们会构建两个比较大的 `Dense` 层，每个层具有 4096 个单元。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0DQM24hL_14Q"
      },
      "outputs": [],
      "source": [
        "inputs = keras.Input(shape=(784,), name='digits')\n",
        "if tf.config.list_physical_devices('GPU'):\n",
        "  print('The model will run with 4096 units on a GPU')\n",
        "  num_units = 4096\n",
        "else:\n",
        "  # Use fewer units on CPUs so the model finishes in a reasonable amount of time\n",
        "  print('The model will run with 64 units on a CPU')\n",
        "  num_units = 64\n",
        "dense1 = layers.Dense(num_units, activation='relu', name='dense_1')\n",
        "x = dense1(inputs)\n",
        "dense2 = layers.Dense(num_units, activation='relu', name='dense_2')\n",
        "x = dense2(x)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2dezdcqnOXHk"
      },
      "source": [
        "每个层都有一条策略，默认情况下会使用全局策略。因此，每个 `Dense` 层都具有 `mixed_float16` 策略，这是因为之前已将 `mixed_float16` 设置为全局策略。这样，dense 层就会执行 float16 计算，并使用 float32 变量。为了执行 float16 计算，它们会将输入转换为 float16 类型，因此，输出也是 float16 类型。它们的变量是 float32 类型，在调用层时，它们会将变量转换为 float16 类型，从而避免 dtype 不匹配所引起的错误。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "kC58MzP4PEcC"
      },
      "outputs": [],
      "source": [
        "print('x.dtype: %s' % x.dtype.name)\n",
        "# 'kernel' is dense1's variable\n",
        "print('dense1.kernel.dtype: %s' % dense1.kernel.dtype.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_WAZeqDyqZcb"
      },
      "source": [
        "接下来创建输出预测。通常，您可以按如下方法创建输出预测，但是对于 float16，其结果不一定具有数值稳定性。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ybBq1JDwNIbz"
      },
      "outputs": [],
      "source": [
        "# INCORRECT: softmax and model output will be float16, when it should be float32\n",
        "outputs = layers.Dense(10, activation='softmax', name='predictions')(x)\n",
        "print('Outputs dtype: %s' % outputs.dtype.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "D0gSWxc9NN7q"
      },
      "source": [
        "模型末尾的 softmax 激活值本应为 float32 类型。但由于 dtype 策略是 `mixed_float16`，softmax 激活通常会使用 float16 dtype 进行计算，并且会输出 float16 张量。\n",
        "\n",
        "这一问题可以通过分离 Dense 和 softmax 层，并将 `dtype='float32'` 传递至 softmax 层来解决。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IGqCGn4BsODw"
      },
      "outputs": [],
      "source": [
        "# CORRECT: softmax and model output are float32\n",
        "x = layers.Dense(10, name='dense_logits')(x)\n",
        "outputs = layers.Activation('softmax', dtype='float32', name='predictions')(x)\n",
        "print('Outputs dtype: %s' % outputs.dtype.name)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tUdkY_DHsP8i"
      },
      "source": [
        "将 `dtype='float32'` 传递至 softmax 层的构造函数会将该层的 dtype 策略重写为 `float32` 策略，从而由后者执行计算并保持变量为 float32 类型。同样，我们也可以传递 `dtype=mixed_precision.Policy('float32')`；层始终将 dtype 参数转换为策略。由于 `Activation` 层没有变量，因此会忽略该策略的变量 dtype，但是该策略的计算 dtype 为 float32，因此 softmax 和模型的输出也是 float32。\n",
        "\n",
        "您可以在模型中间添加 float16 类型的 softmax，但模型末尾的 softmax 应为 float32 类型。原因是，如果从 softmax 传递给损失函数的中间张量是 float16 或 bfloat16 类型，则会出现数值问题。\n",
        "\n",
        "如果您认为使用 float16 计算无法获得数值稳定性，则可以通过传递 `dtype='float32'`，将任何层的 dtype 重写为 float32 类型。但通常，只有模型的最后一层才需要这样重写，因为对大多数层来说，`mixed_float16` 和 `mixed_bfloat16` 的精度已经足够。\n",
        "\n",
        "即使模型不以 softmax 结尾，输出也仍是 float32。虽然对这一特定模型来说并非必需，但可以使用以下代码将模型输出转换为 float32 类型："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dzVAoLI56jR8"
      },
      "outputs": [],
      "source": [
        "# The linear activation is an identity function. So this simply casts 'outputs'\n",
        "# to float32. In this particular case, 'outputs' is already float32 so this is a\n",
        "# no-op.\n",
        "outputs = layers.Activation('linear', dtype='float32')(outputs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tpY4ZP7us5hA"
      },
      "source": [
        "接下来完成并编译该模型，然后生成输入数据。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "g4OT3Z6kqYAL"
      },
      "outputs": [],
      "source": [
        "model = keras.Model(inputs=inputs, outputs=outputs)\n",
        "model.compile(loss='sparse_categorical_crossentropy',\n",
        "              optimizer=keras.optimizers.RMSprop(),\n",
        "              metrics=['accuracy'])\n",
        "\n",
        "(x_train, y_train), (x_test, y_test) = keras.datasets.mnist.load_data()\n",
        "x_train = x_train.reshape(60000, 784).astype('float32') / 255\n",
        "x_test = x_test.reshape(10000, 784).astype('float32') / 255"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Sm8FJHegVRN"
      },
      "source": [
        "本示例将输入数据从 int8 强制转换为 float32。我们不转换为 float16 是因为在 CPU 上除以 255 时，float16 的运算速度比 float32 慢。在这种情况下，性能差距可以忽略不计，但一般来说，在 CPU 上执行运算时，数学处理输入应使用 float32 类型。该模型的第一层会将输入转换为 float16，因为每一层都会将浮点输入强制转换为其计算 dtype。\n",
        "\n",
        "检索模型的初始权重。这样可以通过加载权重来从头开始训练。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0UYs-u_DgiA5"
      },
      "outputs": [],
      "source": [
        "initial_weights = model.get_weights()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zlqz6eVKs9aU"
      },
      "source": [
        "## 使用 Model.fit 训练模型\n",
        "\n",
        "接下来训练模型。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hxI7-0ewmC0A"
      },
      "outputs": [],
      "source": [
        "history = model.fit(x_train, y_train,\n",
        "                    batch_size=8192,\n",
        "                    epochs=5,\n",
        "                    validation_split=0.2)\n",
        "test_scores = model.evaluate(x_test, y_test, verbose=2)\n",
        "print('Test loss:', test_scores[0])\n",
        "print('Test accuracy:', test_scores[1])\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MPhJ9OPWt4x5"
      },
      "source": [
        "请注意，模型会在日志中打印每个样本的时间：例如，“4us/sample”。第一个周期 (Epoch) 可能比较慢，因为 TensorFlow 需要花一些时间来优化模型，但之后每个样本的时间会稳定下来。\n",
        "\n",
        "如果在 Colab 中运行本指南中的示例，您可以使用 float32 对比混合精度的性能。为此，请在“Setting the dtype policy”部分将策略从 `mixed_float16` 更改为 `float32`，然后重新运行所有代码单元，直到此代码点。在计算能力至少为 7.0 的 GPU 上，您会发现每个样本的时间大大增加，表明混合精度提升了模型的速度。例如，对于 Titan V GPU，每样本时间可从 4us 增加到 12us。在继续学习本指南之前，请确保将策略改回 `mixed_float16` 并重新运行代码单元。\n",
        "\n",
        "对于很多实际模型，使用混合精度时还可以将批次大小加倍而不会耗尽内存，因为 float16 张量只需要使用 float32 一半的内存。不过，这对本文中所讲的小模型毫无意义，因为您几乎可以使用任何 dtype 来运行该模型，而每个批次可以包含有 60,000 张图片的整个 MNIST 数据集。\n",
        "\n",
        "如果在 TPU 上运行混合精度，您会发现与在 GPU 上运行混合精度相比，性能提升并不明显。这是因为即使默认 dtype 策略为 `float32`，TPU 也会在后台执行一些 bfloat16 运算。对于某些使用 bfloat16 可以获得数值稳定性的运算（如矩阵乘法），TPU 硬件不支持使用 float32。对于此类运算，TPU 后端会自动在内部使用 bfloat16。因此，将 `dtype='float32'` 传递至使用此类运算的层不会产生数值影响，不过，使用 bfloat16 计算来运行此类层也不会产生负面影响。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mNKMXlCvHgHb"
      },
      "source": [
        "## 损失放大\n",
        "\n",
        "损失放大是 `tf.keras.Model.fit` 自动使用 `mixed_float16` 策略执行，从而避免数值下溢的一种技巧。本节介绍损失放大及其行为定制方法。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1xQX62t2ow0g"
      },
      "source": [
        "### 下溢和溢出\n",
        "\n",
        "float16 数据类型的动态范围比 float32 窄。这意味着大于 $65504$ 的数值会因溢出而变为无穷大，小于 $6.0 \\times 10^{-8}$ 的数值则会因下溢而变成零。float32 和 bfloat16 的动态范围要大得多，因此一般不会出现下溢或溢出的问题。\n",
        "\n",
        "例如："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CHmXRb-yRWbE"
      },
      "outputs": [],
      "source": [
        "x = tf.constant(256, dtype='float16')\n",
        "(x ** 2).numpy()  # Overflow"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5unZLhN0RfQM"
      },
      "outputs": [],
      "source": [
        "x = tf.constant(1e-5, dtype='float16')\n",
        "(x ** 2).numpy()  # Underflow"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pUIbhQypRVe_"
      },
      "source": [
        "实际上，float16 也极少出现下溢的情况。此外，在正向传递中出现下溢的情形更是十分罕见。但是，在反向传递中，梯度可能因下溢而变为零。损失放大就是一个防止出现下溢的技巧。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FAL5qij_oNqJ"
      },
      "source": [
        "### 损失放大背景知识\n",
        "\n",
        "损失放大的基本概念非常简单：只需将损失乘以某个大数字即可，如 $1024$。我们将该数字称为*损失标度*。这会将梯度放大 $1024$ 倍，大大降低了发生下溢的几率。计算出最终梯度后，将其除以 $1024$ 即可得到正确值。\n",
        "\n",
        "该过程的伪代码是：\n",
        "\n",
        "```\n",
        "loss_scale = 1024 loss = model(inputs) loss *= loss_scale # We assume `grads` are float32. We do not want to divide float16 gradients grads = compute_gradient(loss, model.trainable_variables) grads /= loss_scale\n",
        "```\n",
        "\n",
        "选择合适的损失标度比较困难。如果损失标度太小，梯度可能仍会因下溢而变为零。如果太大，则会出现相反的问题：梯度可能因溢出而变为无穷大。\n",
        "\n",
        "为了解决这一问题，TensorFlow 会动态确定损失标度，因此，您不必人工选择。如果使用 `tf.keras.Model.fit`，则会自动完成损失放大，您不必做任何额外的工作。下一节会对此进行详细阐述。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qlbxbjxShf0m"
      },
      "source": [
        "### 选择损失标度\n",
        "\n",
        "每一条 dtype 策略都有一个关联的 `tf.mixed_precision.experimental.LossScale` 对象（可选），这代表一个固定或动态损失标度。默认情况下，`mixed_float16` 策略的损失标度是 `tf.mixed_precision.experimental.DynamicLossScale`，它会动态确定损失标度值。其他策略在默认情况下没有损失标度，因为只有在使用 float16 时才需要该值。您可以查询策略的损失标度："
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0iSbassjdldy"
      },
      "outputs": [],
      "source": [
        "loss_scale = policy.loss_scale\n",
        "print('Loss scale: %s' % loss_scale)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l8c0ULUxdu2b"
      },
      "source": [
        "损失标度会打印很多内部状态信息，不过您可以忽略。最重要的部分是 `current_loss_scale`，这里显示的就是损失标度的当前值。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Mr0cAMCHfhlj"
      },
      "source": [
        "在构建 dtype 策略时，您也可以通过传递一个数字来使用静态损失标度。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "4Sa4bGFif9GW"
      },
      "outputs": [],
      "source": [
        "new_policy = mixed_precision.Policy('mixed_float16', loss_scale=1024)\n",
        "print(new_policy.loss_scale)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "alD_0frtgFKK"
      },
      "source": [
        "dtype 策略的构造函数始终会将损失标度转换为 `LossScale` 对象。在本例中，它会转换为 `tf.mixed_precision.experimental.FixedLossScale`，即除 `DynamicLossScale` 之外的唯一 `LossScale` 子类。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Gi1DLrLF7bnr"
      },
      "source": [
        "注：*建议使用动态损失标度，而不要使用任何其他值*。选择固定损失标度非常困难，因为太小则会导致模型无法得到训练，而太大又会导致梯度中出现 Inf（无穷大）或 NaN（无效）值。动态损失标度的值通常非常接近最佳值，因此，您不必做任何修改。目前，动态损失标度的计算速度略低于固定损失标度，但未来会有所改善。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LJ_1W-Axh2bp"
      },
      "source": [
        "就像层一样，每个模型都有一条 dtype 策略。如果存在，模型会使用其策略的损失标度，在 `tf.keras.Model.fit` 方法中应用损失放大。这意味着如果使用 `Model.fit`，则完全不必担心损失放大：默认情况下，`mixed_float16` 策略有一个动态损失标度，而 `Model.fit` 将应用该值。\n",
        "\n",
        "对于自定义训练循环，模型会忽略该策略的损失标度，因此，您必须手动应用该值。下一节会对此进行阐述。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "yqzbn8Ks9Q98"
      },
      "source": [
        "## 使用自定义训练循环训练模型"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CRANRZZ69nA7"
      },
      "source": [
        "目前，您已经使用混合精度通过 `tf.keras.Model.fit` 对 Keras 模型进行了训练。接下来将使用混合精度执行自定义训练循环。如果您还不了解自定义训练循环，请先阅读[自定义训练指南](https://render.githubusercontent.com/tutorials/customization/custom_training_walkthrough.ipynb)。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wXTaM8EEyEuo"
      },
      "source": [
        "使用混合精度运行自定义训练循环需要对使用 float32 运行训练的模型进行两方面的更改：\n",
        "\n",
        "1. 使用混合精度构建模型（已完成）\n",
        "2. 如果使用 `mixed_float16`，则明确使用损失放大。\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M2zpp7_65mTZ"
      },
      "source": [
        "对于步骤 (2)，您将使用 `tf.keras.mixed_precision.experimental.LossScaleOptimizer` 类，其中会封装一个优化器并应用损失放大。该类有两个参数：优化器和损失标度。像下面这样构造一个类即可使用动态损失标度"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ogZN3rIH0vpj"
      },
      "outputs": [],
      "source": [
        "optimizer = keras.optimizers.RMSprop()\n",
        "optimizer = mixed_precision.LossScaleOptimizer(optimizer, loss_scale='dynamic')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8-_ZZVEC2MZQ"
      },
      "source": [
        "传递 `'dynamic'` 等效于传递 `tf.mixed_precision.experimental.DynamicLossScale()`。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JZYEr5hA3MXZ"
      },
      "source": [
        "接下来定义损失对象和 `tf.data.Dataset`。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9cE7Mm533hxe"
      },
      "outputs": [],
      "source": [
        "loss_object = tf.keras.losses.SparseCategoricalCrossentropy()\n",
        "train_dataset = (tf.data.Dataset.from_tensor_slices((x_train, y_train))\n",
        "                 .shuffle(10000).batch(8192))\n",
        "test_dataset = tf.data.Dataset.from_tensor_slices((x_test, y_test)).batch(8192)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4W0zxrxC3nww"
      },
      "source": [
        "接下来定义训练步骤函数。为了放大损失和缩小梯度，将使用损失标度优化器的两个新方法：\n",
        "\n",
        "- `get_scaled_loss(loss)`：将损失值乘以损失标度值\n",
        "- `get_unscaled_gradients(gradients)`：获取一系列放大的梯度作为输入，并将每一个梯度除以损失标度，从而将其缩小为实际值\n",
        "\n",
        "为了防止梯度发生下溢，必须使用这些函数。随后，如果全部没有出现 Inf 或 NaN 值，则 `LossScaleOptimizer.apply_gradients` 会应用这些梯度。它还会更新损失标度，如果梯度出现 Inf 或 NaN 值，则会将其减半，而如果出现零值，则会增大损失标度。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "V0vHlust4Rug"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def train_step(x, y):\n",
        "  with tf.GradientTape() as tape:\n",
        "    predictions = model(x)\n",
        "    loss = loss_object(y, predictions)\n",
        "    scaled_loss = optimizer.get_scaled_loss(loss)\n",
        "  scaled_gradients = tape.gradient(scaled_loss, model.trainable_variables)\n",
        "  gradients = optimizer.get_unscaled_gradients(scaled_gradients)\n",
        "  optimizer.apply_gradients(zip(gradients, model.trainable_variables))\n",
        "  return loss"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rcFxEjia6YPQ"
      },
      "source": [
        "在训练的开始阶段，`LossScaleOptimizer` 可能会跳过前几个步骤。先使用非常大的损失标度，以便快速确定最佳值。经过几个步骤后，损失标度将稳定下来，这时跳过的步骤将会很少。这一过程是自动执行的，不会影响训练质量。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IHIvKKhg4Y-G"
      },
      "source": [
        "现在定义测试步骤。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "nyk_xiZf42Tt"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def test_step(x):\n",
        "  return model(x, training=False)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hBs98MZyhBOB"
      },
      "source": [
        "加载模型的初始权重，以便从头开始重新训练。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jpzOe3WEhFUJ"
      },
      "outputs": [],
      "source": [
        "model.set_weights(initial_weights)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s9Pi1ADM47Ud"
      },
      "source": [
        "最后，运行自定义训练循环。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "N274tJ3e4_6t"
      },
      "outputs": [],
      "source": [
        "for epoch in range(5):\n",
        "  epoch_loss_avg = tf.keras.metrics.Mean()\n",
        "  test_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(\n",
        "      name='test_accuracy')\n",
        "  for x, y in train_dataset:\n",
        "    loss = train_step(x, y)\n",
        "    epoch_loss_avg(loss)\n",
        "  for x, y in test_dataset:\n",
        "    predictions = test_step(x)\n",
        "    test_accuracy.update_state(y, predictions)\n",
        "  print('Epoch {}: loss={}, test accuracy={}'.format(epoch, epoch_loss_avg.result(), test_accuracy.result()))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "d7daQKGerOFE"
      },
      "source": [
        "## GPU 性能提示\n",
        "\n",
        "下面是在 GPU 上使用混合精度时的一些性能提示。\n",
        "\n",
        "### 增大批次大小\n",
        "\n",
        "当使用混合精度时，如果不影响模型质量，可以尝试使用双倍批次大小运行。因为 float16 张量只使用一半内存，所以，您通常可以将批次大小增大一倍，而不会耗尽内存。增大批次大小通常可以提高训练吞吐量，即模型每秒可以运行的训练元素数量。\n",
        "\n",
        "### 确保使用 GPU Tensor 核心\n",
        "\n",
        "如前所述，现代 NVIDIA GPU 使用称为 Tensor 核心的特殊硬件单元， 可以非常快速地执行 float16 矩阵乘法运算。但是，Tensor 核心要求张量的某些维度是 8 的倍数。在下面的示例中，当且仅当参数值是 8 的倍数时，才能使用 Tensor 核心。\n",
        "\n",
        "- tf.keras.layers.Dense(**units=64**)\n",
        "- tf.keras.layers.Conv2d(**filters=48**, kernel_size=7, stride=3)\n",
        "    - 其他卷积层也是如此，如 tf.keras.layers.Conv3d\n",
        "- tf.keras.layers.LSTM(**units=64**)\n",
        "    - 其他 RNN 也是如此，如 tf.keras.layers.GRU\n",
        "- tf.keras.Model.fit(epochs=2, **batch_size=128**)\n",
        "\n",
        "您应该尽可能使用 Tensor 核心。如果要了解更多信息，请参阅 [NVIDIA 深度学习性能指南](https://docs.nvidia.com/deeplearning/sdk/dl-performance-guide/index.html)，其中介绍了使用 Tensor 核心的具体要求以及与 Tensor 核心相关的其他性能信息。\n",
        "\n",
        "### XLA\n",
        "\n",
        "XLA 是一款可以进一步提高混合精度性能，也可以在较小程度上提高 float32 性能的编译器。有关详细信息，请参阅 [XLA 指南](https://tensorflow.google.cn/xla)。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2tFDX8fm6o_3"
      },
      "source": [
        "## Cloud TPU 性能提示\n",
        "\n",
        "就像在 GPU 上一样，您也应该尝试将批次大小增大一倍，因为 bfloat16 张量同样只使用一半内存。双倍批次大小可能会提高训练吞吐量。\n",
        "\n",
        "TPU 不需要任何其他特定于混合精度的调整即可获得最佳性能。TPU 已经要求使用 XLA。它们可以从某些是 $128$ 的倍数的维度获得优势，不过就像使用混合精度一样，这同样适用于 float32。有关一般 TPU 性能提示，请参阅[Cloud TPU 性能指南](https://cloud.google.com/tpu/docs/performance-guide)。这些提示对混合精度和 float32 均适用。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "--wSEU91wO9w"
      },
      "source": [
        "## 总结\n",
        "\n",
        "- 如果您使用的是计算能力至少为 7.0 的 TPU 或 NVIDIA GPU，则应使用混合精度，因为它可以将性能提升多达 3 倍。\n",
        "- 您可以按如下代码使用混合精度：\n",
        "    ```\n",
        "    # On TPUs, use 'mixed_bfloat16' instead policy = tf.keras.mixed_precision.experimental.Policy('mixed_float16') mixed_precision.set_policy(policy)\n",
        "    ```\n",
        "- 如果您的模型以 softmax 结尾，请确保其类型为 float32。不管您的模型以什么结尾，必须确保输出为 float32。\n",
        "- 如果您通过 `mixed_float16` 使用自定义训练循环，则除了上述几行代码外，您还需要使用 `tf.keras.mixed_precision.experimental.LossScaleOptimizer` 封装您的优化器。然后调用 `optimizer.get_scaled_loss` 来放大损失，并且调用 `optimizer.get_unscaled_gradients` 来缩小梯度。\n",
        "- 如果不会降低计算准确率，则可以将训练批次大小加倍。\n",
        "- 在 GPU 上，确保大部分张量维度是 $8$ 的倍数，从而最大限度提高性能\n",
        "\n",
        "有关使用 `tf.keras.mixed_precision` API 的混合精度的更多示例，请参阅[官方模型仓库](https://github.com/tensorflow/models/tree/master/official)。大多数官方模型（如 [ResNet](https://github.com/tensorflow/models/tree/master/official/vision/image_classification) 和 [Transformer](https://github.com/tensorflow/models/blob/master/official/nlp/transformer)）通过传递 `--dtype=fp16` 来使用混合精度。\n"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "mixed_precision.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
